{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "epochs = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 6 - 使用CNN在MNIST数据集上进行联邦学习\n",
    "\n",
    "## 使用PyTorch + PySyft 10行代码升级到联邦学习\n",
    "\n",
    "\n",
    "### 背景 \n",
    "\n",
    "联邦学习是一种非常令人兴奋且令人振奋的机器学习技术，旨在建立可在分散数据上学习的系统。想法是，数据保留在其生产者（也称为_worker_）的手中，这有助于改善隐私和所有权，并且该模型在工作机之间共享。例如，一种直接的应用程序是在编写文本时预测手机上的下一个单词：您不希望将用于训练的数据（即，短信）发送到中央服务器。\n",
    "\n",
    "因此，联合学习的兴起与数据隐私意识的传播紧密相关，并且自2018年5月起实施数据保护的欧盟GDPR成为催化剂。为了遵循法规，苹果或谷歌等大型参与者已开始对该技术进行大量投资，特别是为了保护移动用户的隐私，但他们尚未提供其工具。在OpenMined，我们相信愿意进行机器学习项目的任何人都应该能够毫不费力地实现隐私保护工具。我们已经构建了用于单行加密数据的工具[如我们的博客文章所述](https://blog.openmined.org/training-cnns-using-spdz/)，现在我们发布了利用新的PyTorch 1.0版本提供了直观的界面来构建安全且可扩展的模型。\n",
    "\n",
    "在这个教程中，我们直接使用了例子 [the canonical example of training a CNN on MNIST using PyTorch](https://github.com/pytorch/examples/blob/master/mnist/main.py) ，展示使用我们的库 [PySyft library](https://github.com/OpenMined/PySyft/)升级为联邦学习是多么简单。我们将遍历示例的每个部分，并在不同的代码后着重标注。\n",
    "\n",
    "你也可以在这里找到这份材料：[our blogpost](https://blog.openmined.org/upgrade-to-federated-learning-in-10-lines)。\n",
    "\n",
    "作者:\n",
    "- Théo Ryffel - GitHub: [@LaRiffle](https://github.com/LaRiffle)\n",
    "\n",
    "中文版译者：\n",
    "- Hou Wei - github：[@dljgs1](https://github.com/dljgs1)\n",
    "\n",
    "**好，让我们开始吧！**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 导包以及模型规格\n",
    "\n",
    "首先导入所需的官方包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后是PySyft多出的部分，特别是要定义远程工作机`alice` 和 `bob`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import syft as sy  # <-- NEW: import the Pysyft library\n",
    "hook = sy.TorchHook(torch)  # <-- NEW: hook PyTorch ie add extra functionalities to support Federated Learning\n",
    "bob = sy.VirtualWorker(hook, id=\"bob\")  # <-- NEW: define remote worker bob\n",
    "alice = sy.VirtualWorker(hook, id=\"alice\")  # <-- NEW: and alice"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们定义学习任务的设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Arguments():\n",
    "    def __init__(self):\n",
    "        self.batch_size = 64\n",
    "        self.test_batch_size = 1000\n",
    "        self.epochs = epochs\n",
    "        self.lr = 0.01\n",
    "        self.momentum = 0.5\n",
    "        self.no_cuda = False\n",
    "        self.seed = 1\n",
    "        self.log_interval = 30\n",
    "        self.save_model = False\n",
    "\n",
    "args = Arguments()\n",
    "\n",
    "use_cuda = not args.no_cuda and torch.cuda.is_available()\n",
    "\n",
    "torch.manual_seed(args.seed)\n",
    "\n",
    "device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "\n",
    "kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 数据加载并发送给工作机\n",
    "我们首先加载数据，然后使用`.federate`方法将训练数据集转换为跨工作人员的联合数据集。现在，该联合数据集已提供给FederatedDataLoader。测试数据集保持不变。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "federated_train_loader = sy.FederatedDataLoader( # <-- this is now a FederatedDataLoader \n",
    "    datasets.MNIST('../data', train=True, download=True,\n",
    "                   transform=transforms.Compose([\n",
    "                       transforms.ToTensor(),\n",
    "                       transforms.Normalize((0.1307,), (0.3081,))\n",
    "                   ]))\n",
    "    .federate((bob, alice)), # <-- NEW: we distribute the dataset across all the workers, it's now a FederatedDataset\n",
    "    batch_size=args.batch_size, shuffle=True, **kwargs)\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST('../data', train=False, transform=transforms.Compose([\n",
    "                       transforms.ToTensor(),\n",
    "                       transforms.Normalize((0.1307,), (0.3081,))\n",
    "                   ])),\n",
    "    batch_size=args.test_batch_size, shuffle=True, **kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CNN规格\n",
    "在这里，我们使用与官方示例中完全相同的CNN。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 20, 5, 1)\n",
    "        self.conv2 = nn.Conv2d(20, 50, 5, 1)\n",
    "        self.fc1 = nn.Linear(4*4*50, 500)\n",
    "        self.fc2 = nn.Linear(500, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.conv1(x))\n",
    "        x = F.max_pool2d(x, 2, 2)\n",
    "        x = F.relu(self.conv2(x))\n",
    "        x = F.max_pool2d(x, 2, 2)\n",
    "        x = x.view(-1, 4*4*50)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = self.fc2(x)\n",
    "        return F.log_softmax(x, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 定义训练和测试函数\n",
    "对于训练功能，由于数据批次分布在`alice`和`bob`之间，因此您需要将模型发送到每个批次的正确位置。 然后，您使用相同的语法远程执行所有操作，就像执行本地PyTorch一样。完成后，您需要恢复模型更新和损失以寻求改进。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(args, model, device, federated_train_loader, optimizer, epoch):\n",
    "    model.train()\n",
    "    for batch_idx, (data, target) in enumerate(federated_train_loader): # <-- now it is a distributed dataset\n",
    "        model.send(data.location) # <-- NEW: send the model to the right location\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = F.nll_loss(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        model.get() # <-- NEW: get the model back\n",
    "        if batch_idx % args.log_interval == 0:\n",
    "            loss = loss.get() # <-- NEW: get the loss back\n",
    "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                epoch, batch_idx * args.batch_size, len(federated_train_loader) * args.batch_size,\n",
    "                100. * batch_idx / len(federated_train_loader), loss.item()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "测试函数完全不变！"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(args, model, device, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            output = model(data)\n",
    "            test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss\n",
    "            pred = output.argmax(1, keepdim=True) # get the index of the max log-probability \n",
    "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "\n",
    "    print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
    "        test_loss, correct, len(test_loader.dataset),\n",
    "        100. * correct / len(test_loader.dataset)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 开始训练！"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "model = Net().to(device)\n",
    "optimizer = optim.SGD(model.parameters(), lr=args.lr) # TODO momentum is not supported at the moment\n",
    "\n",
    "for epoch in range(1, args.epochs + 1):\n",
    "    train(args, model, device, federated_train_loader, optimizer, epoch)\n",
    "    test(args, model, device, test_loader)\n",
    "\n",
    "if (args.save_model):\n",
    "    torch.save(model.state_dict(), \"mnist_cnn.pt\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "等等！ 在这里，您已经使用联邦学习在远程数据上训练了一个模型！"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 最后一件事\n",
    "我知道您很想问一个问题：**与普通的PyTorch相比，进行联邦学习需要多长时间？**\n",
    "\n",
    "实际上，计算时间**少于正常PyTorch执行时间的两倍**！更准确地说，它需要1.9倍的时间，与我们能够添加的功能相比，这几乎是很少的。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 恭喜!!! 是时候加入社区了!\n",
    "\n",
    "祝贺您完成本笔记本教程！ 如果您喜欢此方法，并希望加入保护隐私、去中心化AI和AI供应链（数据）所有权的运动，则可以通过以下方式做到这一点！\n",
    "\n",
    "### 给 PySyft 加星\n",
    "\n",
    "帮助我们的社区的最简单方法是仅通过给GitHub存储库加注星标！ 这有助于提高人们对我们正在构建的出色工具的认识。\n",
    "\n",
    "- [Star PySyft](https://github.com/OpenMined/PySyft)\n",
    "\n",
    "### 选择我们的教程\n",
    "\n",
    "我们编写了非常不错的教程，以更好地了解联合学习和隐私保护学习的外观，以及我们如何为实现这一目标添砖加瓦。\n",
    "\n",
    "- [Checkout the PySyft tutorials](https://github.com/OpenMined/PySyft/tree/master/examples/tutorials)\n",
    "\n",
    "\n",
    "### 加入我们的 Slack!\n",
    "\n",
    "保持最新进展的最佳方法是加入我们的社区！ 您可以通过填写以下表格来做到这一点[http://slack.openmined.org](http://slack.openmined.org)\n",
    "\n",
    "### 加入代码项目!\n",
    "\n",
    "对我们的社区做出贡献的最好方法是成为代码贡献者！ 您随时可以转到PySyft GitHub的Issue页面并过滤“projects”。这将向您显示所有概述，选择您可以加入的项目！如果您不想加入项目，但是想做一些编码，则还可以通过搜索标记为“good first issue”的GitHub问题来寻找更多的“一次性”微型项目。\n",
    "\n",
    "- [PySyft Projects](https://github.com/OpenMined/PySyft/issues?q=is%3Aopen+is%3Aissue+label%3AProject)\n",
    "- [Good First Issue Tickets](https://github.com/OpenMined/PySyft/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22)\n",
    "\n",
    "### 捐赠\n",
    "\n",
    "如果您没有时间为我们的代码库做贡献，但仍想提供支持，那么您也可以成为Open Collective的支持者。所有捐款都将用于我们的网络托管和其他社区支出，例如黑客马拉松和聚会！\n",
    "\n",
    "[OpenMined's Open Collective Page](https://opencollective.com/openmined)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
